#include "SSAORenderPass.h"
#include <SERender/Renderer/Renderer.h>

#include <SERender/Resourses/Loader/TextureLoader.h>
#include <SEngine/Render/Descriptors/SceneDescriptor.h>
#include <SERender/Descriptors/FrameDescriptor.h>
#include <SERender/Renderer/Renderer.h>
#include <SECore/Ecs/Scene.h>
#include <SEngine/Ecs/Components/ModelRenderComp.h>
#include <SEngine/Render/RenderFeature/ModelRenderFeature.h>
#include <SERender/Resourses/Loader/ShaderLoader.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>
#include <SEngine/Render/RenderFeature/UtilRenderFeature.h>
#include <SECore/Manager.h>
#include <random>

#include <SEUi/SEUi.h>
#include <SEUi/Util/Util.h>

namespace SEngine
{    
    SSAORenderPass::SSAORenderPass(Renderer &renderer)
        : RenderPass(renderer)
    {
        m_fbo = std::make_unique<FrameBuffer>();

        TextureLoader<Texture2D> texLoader;
        texLoader.SetInternalFormat(RGB32F);
        texLoader.SetFormat(RGB);
        texLoader.SetType(FLOAT);
        texLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
        texLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
        texLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);
        m_positionRT = texLoader.CreateTexture(1024, 1024, nullptr);
        m_normalRT = texLoader.CreateTexture(1024, 1024, nullptr);
        m_tangentRT = texLoader.CreateTexture(1024, 1024, nullptr);
        texLoader.SetInternalFormat(RGBA32F);
        texLoader.SetFormat(RGBA);
        texLoader.SetType(FLOAT);
        m_SSAOTex = texLoader.CreateTexture(1024, 1024, nullptr);
        m_SSAONoise = CreateNoiseMap(4);

        m_RTMaterial = Manager::GetInstance().Get<Material>("CaptureSSAO").value();
        m_SSAOMaterial = Manager::GetInstance().Get<Material>("SSAO").value();
        m_BlurMaterial = Manager::GetInstance().Get<Material>("Blur").value();
        m_BlurMaterial.Value().GetPSO().blend = true;

        m_SSAOSamples = RandomKernel(64);
    }

    void SSAORenderPass::BeginFrame()
    {
    }

    void SSAORenderPass::DrawFrame()
    {
        auto & modelRenderFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        auto & utilRenderFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UtilRenderFeature>()->get();
        auto & scene = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get().scene;
        auto & frame = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        
        m_renderer.UseFrameBuffer(*m_fbo, [&]() {
            m_fbo->ClearAttachments();
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_positionRT);
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT1, *m_normalRT);
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT2, *m_tangentRT);
            m_fbo->Resize(frame.w, frame.h);
            m_fbo->Use();
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
            scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp& model) {
                modelRenderFeature.DrawModel(model.m_model, m_RTMaterial);
            });
        });

        static float radius = 1.0f;
        static int ssaoKernelSize = 64;
        static int blurKernelSize = 2;
        // ImGui::Begin("RT");
        //     ImGui::DragFloat("Radius", &radius, 0.1f, 0.0f, FLT_MAX);
        //     ImGui::DragInt("SsaoKernelSize", &ssaoKernelSize, 1, 0, 64);
        //     ImGui::DragInt("BlurKernelSize", &blurKernelSize, 1, 1, 4);
        //     if (ImGui::Button("Reload")) {
        //         ResValue(m_RTMaterial).GetShader().lock()->Reload([](auto & path, auto & name) {
        //             ShaderLoader loader;
        //             return loader.Load(path).CreateShader(EngineUniformBlockBinder());
        //         });
        //         ResValue(m_SSAOMaterial).GetShader().lock()->Reload([](auto & path, auto & name) {
        //             ShaderLoader loader;
        //             return loader.Load(path).CreateShader(EngineUniformBlockBinder());
        //         });
        //         ResValue(m_BlurMaterial).GetShader().lock()->Reload([](auto & path, auto & name) {
        //             ShaderLoader loader;
        //             return loader.Load(path).CreateShader(EngineUniformBlockBinder());
        //         });
        //     }
        // ImGui::End();

        m_renderer.UseFrameBuffer(*m_fbo, [&]() {
            m_fbo->ClearAttachments();
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_SSAOTex);
            m_fbo->Resize(frame.w, frame.h);
            m_fbo->Use();
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
            m_SSAOMaterial.Value().Set("u_radius", radius);
            m_SSAOMaterial.Value().Set("u_positionRT", Shader::Sampler{*m_positionRT, TextureSlot::TEXTURE0});
            m_SSAOMaterial.Value().Set("u_normalRT", Shader::Sampler{*m_normalRT, TextureSlot::TEXTURE1});
            m_SSAOMaterial.Value().Set("u_tangentRT", Shader::Sampler{*m_tangentRT, TextureSlot::TEXTURE2});
            m_SSAOMaterial.Value().Set("u_ssaoNoise", Shader::Sampler{*m_SSAONoise, TextureSlot::TEXTURE3});
            for (int i = 0; i < m_SSAOSamples.size(); ++i) {
                m_SSAOMaterial.Value().Set(("u_samples[" + std::to_string(i) + "]").c_str(), m_SSAOSamples[i]);
            }
            m_SSAOMaterial.Value().Set("u_kernelSize", ssaoKernelSize);
            utilRenderFeature.DrawBillboard(m_SSAOMaterial);
        });
        m_BlurMaterial.Value().Set("u_texture", Shader::Sampler{*m_SSAOTex, TextureSlot::TEXTURE0});
        m_BlurMaterial.Value().Set("u_kernelSize", blurKernelSize);
        utilRenderFeature.DrawBillboard(m_BlurMaterial);

        // ShowRT("positionRT", m_positionRT->GetId());
    }
    
    std::vector<glm::vec3> SSAORenderPass::RandomKernel(size_t num)
    {
        std::vector<glm::vec3> ssaoKernel;
        std::uniform_real_distribution<float> randomFloats(0.0, 1.0); // 随机浮点数，范围0.0 - 1.0
        std::default_random_engine generator;
        for (size_t i = 0; i < num; ++i)
        {
            glm::vec3 sample(
                randomFloats(generator) * 2.0 - 1.0, 
                randomFloats(generator) * 2.0 - 1.0, 
                randomFloats(generator)
            );
            sample = glm::normalize(sample);
            sample *= randomFloats(generator);
            float scale = i / 64.0f; 
            scale = glm::mix(0.1f, 1.0f, scale * scale);
            sample *= scale;
            ssaoKernel.push_back(sample);
        }
        return ssaoKernel;
    }

    std::shared_ptr<Texture2D> SSAORenderPass::CreateNoiseMap(int size)
    {
        TextureLoader<Texture2D> texLoader;
        texLoader.SetInternalFormat(RGB16F);
        texLoader.SetFormat(RGB);
        texLoader.SetType(FLOAT);
        texLoader.SetParameter(TEXTURE_MAG_FILTER, NEAREST);
        texLoader.SetParameter(TEXTURE_MIN_FILTER, NEAREST);
        texLoader.SetParameter(TEXTURE_WRAP_S, REPEAT);
        texLoader.SetParameter(TEXTURE_WRAP_S, REPEAT);
        std::uniform_real_distribution<float> randomFloats(0.0, 1.0); // 随机浮点数，范围0.0 - 1.0
        std::default_random_engine generator;
        std::vector<glm::vec3> ssaoNosise;
        for (int i = 0; i < size * size; ++i) {
            ssaoNosise.emplace_back(
                glm::vec3{
                    randomFloats(generator) * 2.0f - 1.0f,
                    randomFloats(generator) * 2.0f - 1.0f,
                    0.0f
                }
            );
        }
        return texLoader.CreateTexture(size, size, &ssaoNosise.data()->x);
    }
} // namespace SEngine
